from __future__ import division
from math import *
from Init_Cond import *
from numpy import *
import pylab as pl
import Laplace
Laplace.f()
import os.path
import Directory
directory = os.getcwd()
import Save
import Plot
import Build
# import CompileCollisions
import Visc3
# #import ViscFlux
import Update
import Restructure
import Timestep
# import Impact
# import dM_dt
import Satellites4
import sys
import Restart

a_Embryo = 1.5*AU
t = 0.0
L = L_Mars
M = M_Mars
R_Embryo = (3.0*M/(4.0*pi*rho))**(1.0/3.0)
v_K = (G*M_Sun/a_Embryo)**(1/2)

inside = 0
Build.f(1, M)

Roche = 2.456*R_Embryo*(rho/1.876)**(1./3.) #denisty of Phobos
i_roche = int(round((Roche - r_I)/deltar - .5))
Rigid_Roche = 1.26*R_Embryo*(rho/1.876)**(1./3.)
i_rigid = int(round((Rigid_Roche - r_I)/deltar - 0.5))

Satellites4.construct(i_rigid, M)
last = len(Sat_M)-1   #location of last satellite (w/mass).  initially set to N to run the whole disk.  $169 for beyond synchro
for a in range(i_roche, N):   #Truncates disk to the roche limit
    m[a] = 0.0
    sigma[a] = m[a]/deltaA[a]

# ###*~*~~**~ If you want to restart and pick up where you left off.  f(restart yes or no, M, are there satellites?)
# Restart.f(1, M, 1)
# import os.path
# import Directory
# directory = os.getcwd()
print(sum(m), sum(Sat_M))
deltat = deltaT
system_mass = sum(m)

end_sim = 5.0000001e9*year
switch = 0

while t <= end_sim:
    Mars_Accrete = 0.0
    T = 4.0*pi*M*R_Embryo**2.0/(5.0*L)
    W = 2.0*pi/T
    r_sync = (T*sqrt(G*M)/(2.0*pi))**(2.0/3.0)
    Roche = 2.456*R_Embryo*(rho/1.876)**(1./3.)
    if t == 0.0:
        Save.f(t, M, Rigid_Roche, Roche, r_sync)


        Sat_r[5] = 1.0
        Sat_M[5] = 1.0
        Plot.sat(t, M, R_Embryo, r_sync, Roche, last, Rigid_Roche)
        Sat_r[5] = Sat_r_default[5]
        Sat_M[5] = 0.0

    #Ring Ring viscosity model
    Visc3.f(1, M, t)   #calculate viscosity values
    Timestep.f(M)    #determine stable timestep
    deltat = Timestep.o(deltat)
    Visc3.g(inside, deltat, M) #calculate sigma
    Mars_Accrete = Visc3.o(Mars_Accrete, inside)
    Visc3.s()

    #Run restructure to conserve angular momentum of ring due to increase mass of the primary
    Restructure.f(M, inside, R_Embryo, W)
    Mars_Accrete = Restructure.o(Mars_Accrete, inside)
    Restructure.g(inside)

    #Update size of embryo. Remove any ring mass swallowed by increase in embyro size.
    Mars_Accrete, inside = Update.f(M, Mars_Accrete, inside, R_Embryo)
    Update.g(inside)

    #Run Satellite evolution
    Satellites4.grow(i_roche, M, deltat, t)
    Satellites4.merge(M, last)
    Satellites4.evolve(t, M, deltat, R_Embryo, r_sync, i_rigid, i_roche, last)
    last = Satellites4.final()

    R_Embryo = (3.0*M/(4.0*pi*rho))**(1.0/3.0)

    #print data every 1000 years.
    if (t/year - (t/year)%1.0)%1.e3 == 0.0:
        plotfile = os.path.join(directory, 'Results(t_' + str(t/year - (t/year)%1.0) + ')' + '.txt')

        if os.path.exists(plotfile):
            pass
        else:
            Save.f(t, M, Rigid_Roche, Roche, r_sync)
            if sum(Sat_M) > 0.0:
                Plot.sat(t, M, R_Embryo, r_sync, Roche, last, Rigid_Roche)


    if t + deltat >= end_sim:
        Plot.sat(t, M, R_Embryo, r_sync, Roche, last)
        # Plot.sats(t, R_Embryo, r_sync, r[i_roche])
        Save.g(t, last)

    if Sat_r[len(Sat_r)-1] > 7*R_Mars:
        donezo = open('final.txt', 'w')
        print(t/year, file = donezo)
        donezo.close()
        Plot.sat(t, M, R_Embryo, r_sync, Roche, last)
        # Plot.sats(t, R_Embryo, r_sync, r[i_roche])
        Save.g(t, last)
        sys.exit()

    t = t + deltat

print(t/year)
